#include "postprocess/decode/lattice_net_decode.h"

#include <numeric>

#include "postprocess/decode/cc_decode.h"

LatticeNetDecode::LatticeNetDecode() noexcept {
  std::cout << "LatticeNetDecode construction function" << std::endl;
  obj_scores_ptr = std::make_shared<std::array<float, kMaxResultSize>>();
  obj_total_scores_ptr =
      std::make_shared<std::array<float, kMaxResultSize * 10>>();
  obj_labels_ptr = std::make_shared<std::array<int32_t, kMaxResultSize>>();
  obj_boxes_ptr =
      std::make_shared<std::array<float, kMaxResultSize * kMaxObjInfoSize>>();
  obj_nms_keep_ptr = std::make_shared<std::array<bool, kMaxResultSize>>();
  map_pts_ptr_ = std::make_shared<
      std::array<float, kMaxResultSize * kMaxMapFeatureSize>>();

  // // object decode parameter
  // std::vector<float> netRange{-51.2, -51.2, -5, 51.2, 51.2, 3.0};
  // std::vector<float> voxelSize{0.2, 0.2, 8};
  // object_decode_para_.bbox_info_number = 9;
  // object_decode_para_.down_size = 4;
  // object_decode_para_.net_range = netRange;
  // object_decode_para_.voxel_size = voxelSize;
  // object_decode_para_.nms_iou_threshold = 0.01;

  // // map decode parameter
  // map_decode_para_.map_info_number = 6;
  // map_decode_para_.preds_number = 2;
  // map_decode_para_.down_size = 4;
  // map_decode_para_.net_range =
  //     std::vector<float>{-51.2, -51.2, -5, 51.2, 51.2, 3.0};
  // map_decode_para_.voxel_size = std::vector<float>{0.2, 0.2, 2};
  // map_decode_para_.type_offset = 3;
  // map_decode_para_.cls_preds_size = std::vector<std::uint32_t>{3, 5, 3, 3};
  // map_decode_para_.start_cls_threshold =
  //     std::vector<float>{0.55, 0.55, 0.55, 0.55};
  // map_decode_para_.cls_threshold = std::vector<float>{0.4, 0.4, 0.4, 0.4};
  // map_decode_para_.vec_threshold = std::vector<float>{45, 45, 45, 45};
  // map_decode_para_.dis_threshold = std::vector<float>{0.4, 0.4, 0.4, 0.4};
  // map_decode_para_.dis_loss_weight = 0.5;
}
LatticeNetDecode::~LatticeNetDecode() {}

bool LatticeNetDecode::DecodeObject(
    autodrive::perception::IONode* dataRaw,
    ObjDecodeParameter const& parameters,
    autodrive::perception::fixed_array<BoundingBox, kMaxBoundingBoxNum>&
        boundingBoxes,
    std::string& err) noexcept {
  std::string errorMsg;
  boundingBoxes.clear();

  // box 0
  bool bRet{decode_box(dataRaw, 10, 0, parameters, boundingBoxes, errorMsg)};
  if (bRet != true) {
    err = "decode box_0 error, " + errorMsg;
    return false;
  }
  return true;
}

bool LatticeNetDecode::decode_box(
    autodrive::perception::IONode* dataRaw, int const& class_number,
    int const& class_base, ObjDecodeParameter const& parameters,
    autodrive::perception::fixed_array<BoundingBox, kMaxBoundingBoxNum>&
        bounding_boxes,
    std::string& err) noexcept {
  auto const& boxInfoNumber{parameters.bbox_info_number};
  auto const& netRange{parameters.net_range};
  auto const& voxelSize{parameters.voxel_size};
  auto const& downSize{parameters.down_size};
  int netLength =
      int((netRange[3] - netRange[0]) / voxelSize[0] / downSize + 0.5);
  int netWidth =
      int((netRange[4] - netRange[1]) / voxelSize[1] / downSize + 0.5);

  // std::array<float, 10> classThreshold{0.4, 0.4, 0.4, 0.4, 0.4,
  //                                      0.4, 0.4, 0.4, 0.4, 0.4};

  char* out_cls_node = static_cast<char*>(dataRaw[0].host_addr);
  char* out_xy_node = static_cast<char*>(dataRaw[5].host_addr);
  char* out_z_node = static_cast<char*>(dataRaw[4].host_addr);
  char* out_lwh_node = static_cast<char*>(dataRaw[3].host_addr);
  char* out_rot_node = static_cast<char*>(dataRaw[2].host_addr);
  char* out_vel_node = static_cast<char*>(dataRaw[1].host_addr);

  std::vector<std::int32_t> out_cls_node_para{
      dataRaw[0].stride, dataRaw[0].shape.at(1), dataRaw[0].shape.at(2),
      dataRaw[0].shape.at(3)};
  std::vector<std::int32_t> out_xy_node_para{
      dataRaw[5].stride, dataRaw[5].shape.at(1), dataRaw[5].shape.at(2),
      dataRaw[5].shape.at(3)};
  std::vector<std::int32_t> out_z_node_para{
      dataRaw[4].stride, dataRaw[4].shape.at(1), dataRaw[4].shape.at(2),
      dataRaw[4].shape.at(3)};
  std::vector<std::int32_t> out_lwh_node_para{
      dataRaw[3].stride, dataRaw[3].shape.at(1), dataRaw[3].shape.at(2),
      dataRaw[3].shape.at(3)};
  std::vector<std::int32_t> out_rot_node_para{
      dataRaw[2].stride, dataRaw[2].shape.at(1), dataRaw[2].shape.at(2),
      dataRaw[2].shape.at(3)};
  std::vector<std::int32_t> out_vel_node_para{
      dataRaw[1].stride, dataRaw[1].shape.at(1), dataRaw[1].shape.at(2),
      dataRaw[1].shape.at(3)};

  std::int8_t scales[6] = {
      dataRaw[0].scaler_bit, dataRaw[5].scaler_bit, dataRaw[4].scaler_bit,
      dataRaw[3].scaler_bit, dataRaw[2].scaler_bit, dataRaw[1].scaler_bit,
  };

  // auto nbDecode{decode_chw_center<decode_type>(
  //     out_cls_node, out_xy_node, out_z_node, out_lwh_node, out_rot_node,
  //     out_vel_node, classThreshold.data(), obj_scores_ptr->data(),
  //     obj_boxes_ptr->data(), obj_labels_ptr->data(),
  //     static_cast<int>(LatticeNetDecode::kMaxResultSize), 0, netLength,
  //     netWidth, class_number, class_base, boxInfoNumber, downSize,
  //     voxelSize.data(), netRange.data(), scales)};

  auto nbDecode{decode_chw_center<decode_type>(
      out_cls_node, out_xy_node, out_z_node, out_lwh_node, out_rot_node,
      out_vel_node, out_cls_node_para.data(), out_xy_node_para.data(),
      out_z_node_para.data(), out_lwh_node_para.data(),
      out_rot_node_para.data(), out_vel_node_para.data(),
      parameters.class_threshold.data(), obj_scores_ptr->data(),
      obj_boxes_ptr->data(), obj_labels_ptr->data(),
      obj_total_scores_ptr->data(),
      static_cast<int>(LatticeNetDecode::kMaxResultSize), 0, netLength,
      netWidth, class_number, class_base, boxInfoNumber, downSize,
      voxelSize.data(), netRange.data(), scales)};
  if (nbDecode == -1) {
    err = "decode object error";
    return -1;
  }

  // nms
  obj_nms_keep_ptr->fill(true);
  nms_cpu(obj_boxes_ptr->data(), obj_scores_ptr->data(),
          obj_nms_keep_ptr->data(), boxInfoNumber, parameters.nms_iou_threshold,
          nbDecode);

  for (int i = 0; i < nbDecode; i++) {
    if (bounding_boxes.full()) {
      break;
    }
    if (obj_nms_keep_ptr->at(i)) {
      auto box{bounding_boxes.end()};
      box->score = obj_scores_ptr->at(i);
      box->label = static_cast<std::uint16_t>(obj_labels_ptr->at(i));
      box->center.x = obj_boxes_ptr->at(i * boxInfoNumber + 0);
      box->center.y = obj_boxes_ptr->at(i * boxInfoNumber + 1);
      box->center.z = obj_boxes_ptr->at(i * boxInfoNumber + 2);
      box->length = obj_boxes_ptr->at(i * boxInfoNumber + 3);
      box->width = obj_boxes_ptr->at(i * boxInfoNumber + 4);
      box->height = obj_boxes_ptr->at(i * boxInfoNumber + 5);
      box->r = obj_boxes_ptr->at(i * boxInfoNumber + 6);
      box->heading = obj_boxes_ptr->at(i * boxInfoNumber + 6);
      box->vx = obj_boxes_ptr->at(i * boxInfoNumber + 7);
      box->vy = obj_boxes_ptr->at(i * boxInfoNumber + 8);
      for (int j = 0; j < 10; ++j) {
        box->scores[j] = obj_total_scores_ptr->at(i * class_number + j);
      }
      gen_polygon(*box, box->polygon);
      bounding_boxes.resize(bounding_boxes.size() + 1);
    }
  }
  return true;
}

void LatticeNetDecode::gen_polygon(
    BoundingBox const& bbox,
    autodrive::perception::fixed_array<autodrive::perception::Point2f, 4>&
        polygon) noexcept {
  Eigen::Vector2f box_centor, dirt;
  Eigen::Vector2f dirt_length;
  Eigen::Vector2f dirt_width;
  box_centor[0] = bbox.center.x;
  box_centor[1] = bbox.center.y;
  dirt[0] = cos(bbox.heading);
  dirt[1] = sin(bbox.heading);
  dirt_length[0] = dirt[0];
  dirt_length[1] = dirt[1];
  dirt_width[0] = -dirt[1];
  dirt_width[1] = dirt[0];

  polygon.clear();
  Eigen::Vector2f p{box_centor +
                    (bbox.length * dirt_length + bbox.width * dirt_width) / 2};
  autodrive::perception::Point2f point0;
  point0.x = p[0];
  point0.y = p[1];
  polygon.add(point0);

  p = box_centor + (bbox.length * dirt_length - bbox.width * dirt_width) / 2;
  autodrive::perception::Point2f point1;
  point1.x = p[0];
  point1.y = p[1];
  polygon.add(point1);

  p = box_centor - (bbox.length * dirt_length + bbox.width * dirt_width) / 2;
  autodrive::perception::Point2f point2;
  point2.x = p[0];
  point2.y = p[1];
  polygon.add(point2);

  p = box_centor - (bbox.length * dirt_length - bbox.width * dirt_width) / 2;
  autodrive::perception::Point2f point3;
  point3.x = p[0];
  point3.y = p[1];
  polygon.add(point3);
}
